#define STB_IMAGE_IMPLEMENTATION
#include "./stb/stb_image.h"
#define STB_IMAGE_WRITE_IMPLEMENTATION
#include "./stb/stb_image_write.h"
#define STB_IMAGE_RESIZE_IMPLEMENTATION
#include "./stb/stb_image_resize.h"
#include "./stb/drawing.hpp"
#include "./stb/common_data_struct.h"
#include <stdint.h>
#include <stdio.h>
#include <string>
#include <cstring>  // 添加memset头文件
#include <cmath>    // 添加fminf头文件
#include <iostream>
#include <stdlib.h>  // 添加 abs() 函数头文件


// 图像基本信息结构体 (兼容原有代码)
typedef struct {
    uint8_t* data;   // 像素数据指针 (RGBA格式)
    int width;       // 图像宽度
    int height;      // 图像高度
    int channels; 
} FrameInfo;

typedef struct {
    float x0, y0, x1, y1;
} Box_xyxy;

typedef struct {
    float x0, y0, w, h;
} Box_xywh;

static const uint8_t color[3] = {0xff, 0, 0};


uint8_t* rgb_to_rgba(const uint8_t* src, int width, int height) {
    if (!src || width <= 0 || height <= 0) return NULL;
    
    uint8_t* dst = (uint8_t*)malloc(width * height * 4);
    if (!dst) return NULL;

    for (int i = 0; i < width * height; i++) {
        dst[i*4]   = src[i*3];     // R
        dst[i*4+1] = src[i*3+1];   // G
        dst[i*4+2] = src[i*3+2];   // B
        dst[i*4+3] = 0;        // A
    }
    return dst;
}

uint8_t* rgba_to_rgb(const uint8_t* src, int width, int height) {
    if (!src || width <= 0 || height <= 0) return NULL;
    
    uint8_t* dst = (uint8_t*)malloc(width * height * 3);
    if (!dst) return NULL;

    for (int i = 0; i < width * height; i++) {
        dst[i*3]   = src[i*4];     // R
        dst[i*3+1] = src[i*4+1];   // G
        dst[i*3+2] = src[i*4+2];   // B
        // 忽略src[i*4+3] (Alpha)
    }
    return dst;
}


// 读取图像并填充结构体，返回是否成功
bool read_img(Img& img, const std::string& img_path) {
    printf("Loading image: %s\n", img_path.c_str());

    // 先清空输出结构体
    memset(&img, 0, sizeof(Img));

    int ori_img_w, ori_img_h, comp;
    unsigned char* imagedata = stbi_load(img_path.c_str(), &ori_img_w, &ori_img_h, &comp, 3);
    if (!imagedata) {
        fprintf(stderr, "Error loading image: %s\n", img_path.c_str());
        return false;
    }
    // 填充Img
    img.w = ori_img_w;
    img.h = ori_img_h;
    img.c = comp;        
    img.w_stride = ori_img_w * comp;
    img.data = imagedata;

    printf("Image loaded successfully: width=%d, height=%d, channels=%d\n", ori_img_w, ori_img_h, comp);
    return true;
}

void image_to_frame(const Img& img, FrameInfo& frame) {
    // 清空输出结构体
    memset(&frame, 0, sizeof(FrameInfo));
    frame.data = rgb_to_rgba(img.data, img.w, img.h);
    frame.width = img.w;
    frame.height = img.h;
    frame.channels = 4;
}


void resize_image(const Img& src_img, Img& dst_img, int dst_width, int dst_height) {
    // 清空输出结构体
    memset(&dst_img, 0, sizeof(Img));
    
    if (!src_img.data || dst_width <= 0 || dst_height <= 0) {
        fprintf(stderr, "Invalid resize parameters\n");
        return;
    }

    // 初始化输出图像结构
    dst_img.w = dst_width;
    dst_img.h = dst_height;
    dst_img.c = src_img.c;
    dst_img.w_stride = dst_width * src_img.c;
    dst_img.data = (uint8_t*)malloc(dst_height * dst_img.w_stride);
    
    if (!dst_img.data) {
        fprintf(stderr, "Memory allocation failed for %d bytes\n", 
                dst_height * dst_img.w_stride);
        return;
    }

    // 调用stb_image_resize进行缩放
    stbir_resize(
        src_img.data, src_img.w, src_img.h, src_img.w_stride,
        dst_img.data, dst_width, dst_height, dst_img.w_stride,
        STBIR_TYPE_UINT8, src_img.c,
        (src_img.c == 4) ? STBIR_ALPHA_CHANNEL_NONE : 0,
        STBIR_FLAG_ALPHA_PREMULTIPLIED,
        STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
        STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT,
        STBIR_COLORSPACE_SRGB, NULL
    );
}

void yolo_resize_img(const Img& src_img, Img& dst_img, int target_size) {
    // 清空输出结构体
    memset(&dst_img, 0, sizeof(Img));
    const uint8_t fill_color[3] = {114, 114, 114}; // YOLO标准填充色
    if (!src_img.data || target_size <= 0) {
        fprintf(stderr, "Invalid resize parameters\n");
        return;
    }

    // 计算缩放比例（保持宽高比）
    float scale = fminf((float)target_size / src_img.w, (float)target_size / src_img.h);
    int new_w = (int)(src_img.w * scale);
    int new_h = (int)(src_img.h * scale);

    // 计算填充位置（居中）
    int dw = target_size - new_w;
    int dh = target_size - new_h;
    int top = dh / 2;
    int bottom = dh - top;
    int left = dw / 2;
    int right = dw - left;

    // 初始化输出图像
    dst_img.w = target_size;
    dst_img.h = target_size;
    dst_img.c = src_img.c;
    dst_img.w_stride = target_size * src_img.c;
    dst_img.data = (uint8_t*)malloc(target_size * target_size * src_img.c);
    
    if (!dst_img.data) {
        fprintf(stderr, "Memory allocation failed\n");
        return;
    }

    // 创建填充画布
    for (int i = 0; i < target_size * target_size; i++) {
        for (int c = 0; c < src_img.c; c++) {
            if (c < 3) { // RGB通道
                dst_img.data[i * src_img.c + c] = color[c];
            } else if (c == 3) { // Alpha通道（如果有）
                dst_img.data[i * src_img.c + c] = 255; // 完全不透明
            }
        }
    }

    // 缩放源图像
    uint8_t* resized_data = (uint8_t*)malloc(new_w * new_h * src_img.c);
    stbir_resize(
        src_img.data, src_img.w, src_img.h, src_img.w_stride,
        resized_data, new_w, new_h, new_w * src_img.c,
        STBIR_TYPE_UINT8, src_img.c,
        (src_img.c == 4) ? STBIR_ALPHA_CHANNEL_NONE : 0,
        STBIR_FLAG_ALPHA_PREMULTIPLIED,
        STBIR_EDGE_CLAMP, STBIR_EDGE_CLAMP,
        STBIR_FILTER_DEFAULT, STBIR_FILTER_DEFAULT,
        STBIR_COLORSPACE_SRGB, NULL
    );

    // 将缩放后的图像复制到画布中央
    for (int y = 0; y < new_h; y++) {
        for (int x = 0; x < new_w; x++) {
            for (int c = 0; c < src_img.c; c++) {
                int dst_pos = ((y + top) * target_size + (x + left)) * src_img.c + c;
                int src_pos = (y * new_w + x) * src_img.c + c;
                dst_img.data[dst_pos] = resized_data[src_pos];
            }
        }
    }

    free(resized_data);
}

/**
 * @brief 裁剪图像（优化版：整行复制）
 * @param in_img  输入图像结构体
 * @param box     裁剪区域 (x0,y0) -> (x1,y1)
 * @param out_img 输出图像结构体（需手动释放 data）
 */
void crop_image(const Img& in_img, const Box_xyxy& box, Img& out_img) {
    // 清空输出结构体
    memset(&out_img, 0, sizeof(Img));
    
    // 1. 边界检查
    int x0 = (int)box.x0;
    int y0 = (int)box.y0;
    int x1 = (int)box.x1;
    int y1 = (int)box.y1;

    // 确保坐标不越界
    x0 = (x0 < 0) ? 0 : (x0 >= in_img.w ? in_img.w - 1 : x0);
    y0 = (y0 < 0) ? 0 : (y0 >= in_img.h ? in_img.h - 1 : y0);
    x1 = (x1 < 0) ? 0 : (x1 >= in_img.w ? in_img.w - 1 : x1);
    y1 = (y1 < 0) ? 0 : (y1 >= in_img.h ? in_img.h - 1 : y1);

    // 确保 x1 > x0 且 y1 > y0
    if (x1 <= x0 || y1 <= y0) {
        fprintf(stderr, "Invalid crop box: x0=%d, y0=%d, x1=%d, y1=%d\n", x0, y0, x1, y1);
        return;
    }

    // 2. 计算裁剪区域
    int crop_w = x1 - x0;
    int crop_h = y1 - y0;
    out_img.w = crop_w;
    out_img.h = crop_h;
    out_img.c = in_img.c;
    out_img.w_stride = crop_w * in_img.c;
    out_img.data = (uint8_t*)malloc(crop_h * out_img.w_stride);

    if (!out_img.data) {
        fprintf(stderr, "Memory allocation failed for %d bytes\n", crop_h * out_img.w_stride);
        return;
    }
    
    // 3. 整行复制（优化关键点）
    for (int y = 0; y < crop_h; ++y) {
        // 计算源图像和目标的起始位置
        int src_pos = (y0 + y) * in_img.w_stride + x0 * in_img.c;
        int dst_pos = y * out_img.w_stride;
        
        // 复制整行数据
        memcpy(out_img.data + dst_pos, in_img.data + src_pos, crop_w * in_img.c);
    }
}

void draw_boxes(Img &img, const Box_xyxy &box) {
    if (!img.data) {
        fprintf(stderr, "Cannot draw on empty image\n");
        return;
    }

    Point pt1 = {.x = (int)box.x0, .y = (int)box.y0};
    Point pt2 = {.x = (int)box.x1, .y = (int)box.y1};
    sample_draw_box_for_image(&img, pt1, pt2, color, 2);
}

// 释放图像内存
void free_img(Img& img) {
    if (img.data) {
        stbi_image_free(img.data);
        img.data = nullptr;
    }
    memset(&img, 0, sizeof(Img));
}

void free_frame(FrameInfo& frame) {
    if (frame.data) {
        stbi_image_free(frame.data);
        frame.data = nullptr;
    }
    memset(&frame, 0, sizeof(FrameInfo));
}

/**
 * @brief RGB图像转灰度图
 * @param rgb_img  输入RGB图像 (3通道)
 * @param gray_img 输出灰度图 (1通道，需手动释放data)
 * @return 是否成功
 * @note 输入必须是3通道（RGB），输出为1通道（灰度）
 */
bool rgb_to_grayscale(const Img& rgb_img, Img& gray_img) {
    // 清空输出结构体
    memset(&gray_img, 0, sizeof(Img));

    // 检查输入有效性
    if (!rgb_img.data || rgb_img.c != 3) {
        fprintf(stderr, "Error: Input must be 3-channel RGB image\n");
        return false;
    }

    // 初始化输出结构
    gray_img.w = rgb_img.w;
    gray_img.h = rgb_img.h;
    gray_img.c = 1;  // 灰度图为1通道
    gray_img.w_stride = rgb_img.w;  // 1通道时，步长=宽度
    gray_img.data = (uint8_t*)malloc(rgb_img.w * rgb_img.h);
    
    if (!gray_img.data) {
        fprintf(stderr, "Error: Memory allocation failed\n");
        return false;
    }

    // 灰度化公式: Gray = 0.299*R + 0.587*G + 0.114*B
    for (int y = 0; y < rgb_img.h; ++y) {
        for (int x = 0; x < rgb_img.w; ++x) {
            int src_pos = y * rgb_img.w_stride + x * 3;
            uint8_t r = rgb_img.data[src_pos];
            uint8_t g = rgb_img.data[src_pos + 1];
            uint8_t b = rgb_img.data[src_pos + 2];
            
            // 整数运算优化（避免浮点计算）
            gray_img.data[y * gray_img.w_stride + x] = 
                (uint8_t)((r * 299 + g * 587 + b * 114 + 500) / 1000);
        }
    }

    return true;
}

#include <stdlib.h>  // 添加 abs() 函数头文件

/**
 * @brief 计算图像差分并统计变化像素
 * @param gray1         输入灰度图1 (1通道)
 * @param gray2         输入灰度图2 (1通道)
 * @param threshold     阈值（0~255）
 * @param total_pixels   [输出] 总像素数
 * @param changed_pixels [输出] 变化像素数
 * @param change_ratio   [输出] 变化比例 (0.0~1.0)
 * @return 是否成功
 */
bool image_absdiff_stats(
    const Img& gray1, const Img& gray2, 
    uint8_t threshold,
    int* total_pixels, int* changed_pixels, float* change_ratio
) {
    // 初始化输出
    *total_pixels = 0;
    *changed_pixels = 0;
    *change_ratio = 0.0f;

    // 检查输入有效性
    if (!gray1.data || !gray2.data || gray1.c != 1 || gray2.c != 1) {
        fprintf(stderr, "Error: Inputs must be 1-channel grayscale images\n");
        return false;
    }
    if (gray1.w != gray2.w || gray1.h != gray2.h) {
        fprintf(stderr, "Error: Image sizes mismatch (%dx%d vs %dx%d)\n",
                gray1.w, gray1.h, gray2.w, gray2.h);
        return false;
    }

    // 计算总像素数
    *total_pixels = gray1.w * gray1.h;

    // 统计变化像素
    for (int i = 0; i < *total_pixels; ++i) {
        int val = abs((int)gray1.data[i] - (int)gray2.data[i]);
        if (val > threshold) {
            (*changed_pixels)++;
        }
    }

    // 计算变化比例
    *change_ratio = (float)(*changed_pixels) / (float)(*total_pixels);

    return true;
}

/**
 * @brief 将xywh格式的框转换为xyxy格式
 * @param box_xywh 输入框 [x0, y0, width, height]
 * @return 输出框 [x0, y0, x1, y1]
 */
Box_xyxy xywh_to_xyxy(const Box_xywh &box_xywh) {
    Box_xyxy box_xyxy;
    box_xyxy.x0 = box_xywh.x0;                    // x0 保持不变
    box_xyxy.y0 = box_xywh.y0;                    // y0 保持不变
    box_xyxy.x1 = box_xywh.x0 + box_xywh.w;      // x1 = x0 + width
    box_xyxy.y1 = box_xywh.y0 + box_xywh.h;      // y1 = y0 + height
    return box_xyxy;
}

int main() {
    Img img1, img2;
    std::string image_path1 = "imgs/cat.png";
    std::string image_path2 = "imgs/cat_shift.png";

    // // 读取图片
    if (!read_img(img1, image_path1)) {
        fprintf(stderr, "Failed to load input image\n");
        return -1;
    }
    if (!read_img(img2, image_path2)) {
        fprintf(stderr, "Failed to load input image\n");
        return -1;
    }
    
    // // 定义裁剪区域
    Box_xywh box_xywh = {350, 420, 600, 900};
    Box_xyxy box_xyxy = xywh_to_xyxy(box_xywh);

    // // 裁剪图片
    Img crop_img1, crop_img2;
    crop_image(img1, box_xyxy, crop_img1);
    crop_image(img2, box_xyxy, crop_img2);
    if (!crop_img1.data || !crop_img2.data) {
        fprintf(stderr, "Failed to crop image\n");
        free_img(img1);
        free_img(img2);
        return -1;
    }


    printf("Cropped image dimensions: %d x %d\n", crop_img1.w, crop_img1.h);
    // // 保存裁剪后的图片
    if (!stbi_write_png("./imgs/cropped.png", crop_img1.w, crop_img1.h, 
        crop_img1.c, crop_img1.data, crop_img1.w_stride)) {
        fprintf(stderr, "Failed to save cropped image\n");
        }    
    printf("Cropped image dimensions: %d x %d\n", crop_img2.w, crop_img2.h);
        // // 保存裁剪后的图片
    if (!stbi_write_png("./imgs/cropped.png", crop_img2.w, crop_img2.h, 
            crop_img2.c, crop_img2.data, crop_img2.w_stride)) {
            fprintf(stderr, "Failed to save cropped image\n");
            }   

    Img gray_img1;
    rgb_to_grayscale(crop_img1, gray_img1);
    Img gray_img2;
    rgb_to_grayscale(crop_img2, gray_img2);
    printf("Grayscale image dimensions: %d x %d\n", gray_img1.w, gray_img1.h);

    int total_pixels, changed_pixels;
    float change_ratio;
    float threshold = 30;
    if (image_absdiff_stats(gray_img1, gray_img2, 30, &total_pixels, &changed_pixels, &change_ratio)) {
        printf("Total pixels: %d\n", total_pixels);
        printf("Changed pixels: %d\n", changed_pixels);
        printf("Change ratio: %.4f\n", change_ratio * 100);
        printf("Threshold: %.4f\n", threshold);
    }

    if (gray_img1.data) {
        if (!stbi_write_png("./imgs/gray1.png", gray_img1.w, gray_img1.h, 
                          gray_img1.c, gray_img1.data, gray_img1.w_stride)) {
            fprintf(stderr, "Failed to save gray image\n");
        }
        free_img(gray_img1);
    }
    if (gray_img2.data) {
        if (!stbi_write_png("./imgs/gray2.png", gray_img2.w, gray_img2.h, 
                          gray_img2.c, gray_img2.data, gray_img2.w_stride)) {
            fprintf(stderr, "Failed to save gray image\n");
        }
        free_img(gray_img2);
    }
   


    // 释放内存
    free_img(img1);
    free_img(img2);
    free_img(crop_img1);
    free_img(crop_img2);
    printf("All operations completed\n");
    return 0;
}